import json
from typing import Optional

from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
from pydantic import BaseModel, Field

from langchain_demo.my_llm import llm

# 生成一个笑话段子，三个属性

# 使用pydantic定义一个类
# class Joke(BaseModel):
#     """笑话搞笑段子的结构类（数据模型类 povo）"""
#     setup: str = Field(description="笑话的开头部分")
#     punchline: str = Field(description="笑话的包袱/笑点")
#     rating: Optional[int] = Field(description="笑话的有趣程度评分，范围1到10")  # 可选的笑话评分字段
#
# # 消息占位符
# prompt_template = ChatPromptTemplate.from_messages(
#     [MessagesPlaceholder('input')]
# )
#
# runnable = llm.with_structured_output(Joke)
# # runnable = llm
#
# chain = prompt_template | runnable
# resp = chain.invoke({"input": [HumanMessage(content="帮我生成一个关于猫的笑话")]})
# print(resp)
#
# print(resp.__dict__)
#
# print(json.dumps(resp.__dict__))


# 使用pydantic定义一个类
class Joke(BaseModel):
    """笑话（搞笑段子）的结构类(数据模型类 POVO)"""
    setup: str = Field(description="笑话的开头部分")  # 笑话的铺垫部分
    punchline: str = Field(description="笑话的包袱/笑点")  # 笑话的爆笑部分
    rating: Optional[int] = Field(description="笑话的有趣程度评分，范围1到10")  # 可选的笑话评分字段


prompt_template = PromptTemplate.from_template("帮我生成一个关于{topic}的笑话。")
runnable = llm.with_structured_output(Joke)

chain = prompt_template | runnable
resp = chain.invoke({"topic": "猫"})
print(resp)

print(resp.__dict__)

json_str = json.dumps(resp.__dict__)
print(json_str)
